package com.nbsaas.boot.controller.stream;

import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@RestController
@RequestMapping("/sse")
public class SseController {

    // 存储每个客户端的SseEmitter连接
    private final ConcurrentHashMap<String, SseEmitter> clients = new ConcurrentHashMap<>();

    // 客户端连接SSE流
    @GetMapping("/connect/{clientId}")
    public SseEmitter connect(@PathVariable String clientId) {
        // 创建SseEmitter
        SseEmitter emitter = new SseEmitter();
        emitter.onTimeout(() -> clients.remove(clientId));  // 超时处理
        emitter.onCompletion(() -> clients.remove(clientId));  // 完成时移除
        // 保存SseEmitter
        clients.put(clientId, emitter);
        return emitter;
    }

    // 向特定客户端发送消息
    @GetMapping("/send/{clientId}")
    public String sendMessage(@PathVariable String clientId, @RequestParam String message) {
        SseEmitter emitter = clients.get(clientId);
        if (emitter != null) {
            try {
                emitter.send(message);
                return "Message sent to client " + clientId;
            } catch (Exception e) {
                return "Failed to send message: " + e.getMessage();
            }
        } else {
            return "Client " + clientId + " not connected.";
        }
    }

    // 向所有连接的客户端广播消息
    @GetMapping("/broadcast")
    public String broadcastMessage(@RequestParam String message) {
        for (SseEmitter emitter : clients.values()) {
            try {
                emitter.send(message);
            } catch (Exception e) {
                // 处理发送错误
            }
        }
        return "Message broadcasted to all clients";
    }
}
